#include "bootinfo.h"
#include <rtl/km/cpu.h>
#include "drivers/usbd.h"
#include "drivers/usbtty.h"
#include "drivers/timer.h"
#include <rtl/print.h>
#include "trustzone.h"

extern UINT32 _bootinfo;
extern UINT32 _bootinfo_length;
extern int g_usb_port_state;

static BOOLEAN usb_port_down(UINT32 tmo_ms)
{
    ULONG start_time = get_timer(0);

    /* check if usb comport close */
    if (!(g_usb_port_state & 1))
        return TRUE;

    while (1)
    {

        //if (get_timer(start_time) > tmo_ms)
        //return FALSE;

        /* kick watchdog to avoid cpu reset */
        //platform_wdt_all_kick();

        mt_usbtty_query_data_size();

        /* check if usb comport close */
        if (!(g_usb_port_state & 1))
            break;

        udelay(20000); /* 20ms */
    }

    return TRUE;
}

extern void apmcu_icache_invalidate();
extern void apmcu_isb();
extern void apmcu_disable_icache();
extern void apmcu_disable_smp();
VOID AgpPrepareToJump()
{
    mt_usbtty_flush();
    usb_port_down(0);
    udc_stop();
    mt_usb_connect_internal();

    apmcu_icache_invalidate();
    apmcu_disable_icache();
    apmcu_isb();
    apmcu_disable_smp();

    // Disable I-Cache D-Cache and branch prediction
    // RtlArmAndSctlr(~(ARM_SCTLR_C | ARM_SCTLR_I | ARM_SCTLR_Z));
}

extern void _jump(UINTPTR Entry, UINTPTR Arg1, UINTPTR Arg2);
VOID AgpBootLk(PVOID Base, SIZE Size, UINTPTR Entry, PVOID AtfBase, SIZE AtfSize)
{
    PATAG_BOOT_MODE BootMode = (PATAG_BOOT_MODE)AgpGetAtag(ATAG_TYPE_BOOT_MODE);
    if (BootMode != NULL)
    {
        BootMode->BootMode = FASTBOOT;
    }
    else
    {
        // TODO: add new ATAG
    }
    PATAG_BOOT_REASON BootReason = (PATAG_BOOT_REASON)AgpGetAtag(ATAG_TYPE_BOOT_REASON);
    if (BootReason)
    {
        BootReason->BootReason = BOOT_REASON_POWER_KEY;
    }

    RtlPrintf("Booting LK at %p entry=%p size=%d arg1=%p arg2=%p ATF base %p size %d\n",
              (PVOID)Base,
              (PVOID)Entry,
              Size,
              (PVOID)_bootinfo,
              (PVOID)_bootinfo_length,
              AtfBase,
              AtfSize);
    if (AtfBase)
    {
        AgpBootAtf(AtfBase, AtfSize, AtfBase, NULL, 0, NULL, Base, Size, (PVOID)Entry);
    }
    else
    {
        AgpPrepareToJump();
        _jump(Entry, _bootinfo, _bootinfo_length);
    }
}

VOID AgpBootPl(UINTPTR Entry)
{
    AgpPrepareToJump();
    _jump(Entry, 0, 0);
}

typedef struct _DA_ARG
{
    UINT32 Magic;
    UINT32 Version;
    UINT32 Flags;
} DA_ARG, *PDA_ARG;

#define DA_ARG_MAGIC 0x58885168
#define DA_ARG_VER 1

#define DA_FLAG_SKIP_PLL_INIT 0x00000001
#define DA_FLAG_SKIP_EMI_INIT 0x00000002

VOID AgpBootMtkDa(UINTPTR DaStart, SIZE DaSize, UINTPTR LkEntry, BOOLEAN UseLk)
{
    DA_ARG _Arg;
    PDA_ARG Arg = &_Arg;
    //PDA_ARG Arg = (PDA_ARG)(ALIGN_UP(DaStart + DaSize, 16));
    Arg->Magic = DA_ARG_MAGIC;
    Arg->Version = DA_ARG_VER;
    // FIXME: currently we assume hardware has already been initialized
    Arg->Flags = DA_FLAG_SKIP_PLL_INIT | DA_FLAG_SKIP_EMI_INIT;

    if (UseLk)
    {
        PATAG_DA_INFO DaInfo = (PATAG_DA_INFO)AgpGetAtag(ATAG_TYPE_DA_INFO);
        if (DaInfo)
        {
            DaInfo->Address = DaStart;
            DaInfo->Size = 0x5A5B5A5B;
            DaInfo->SignatureSize = 0x5A5B5A5B;
            DaInfo->Arg1 = (UINT32)Arg;
            DaInfo->Arg2 = (UINT32)sizeof(DA_ARG);
        }
        else
        {
            // TODO: add new ATAG
        }

        PATAG_BOOT_REASON BootReason = (PATAG_BOOT_REASON)AgpGetAtag(ATAG_TYPE_BOOT_REASON);
        if (BootReason)
        {
            BootReason->BootReason = BOOT_REASON_POWER_KEY;
        }

        PATAG_BOOT_MODE BootMode = (PATAG_BOOT_MODE)AgpGetAtag(ATAG_TYPE_BOOT_MODE);
        if (BootMode != NULL)
        {
            BootMode->BootMode = DOWNLOAD_BOOT;
        }

        AgpPrepareToJump();
        _jump(LkEntry, _bootinfo, _bootinfo_length);
    }
    else
    {
        AgpPrepareToJump();
        _jump(DaStart, (UINTPTR)Arg, sizeof(DA_ARG));
    }
}